# v2.0 包含单文件传输和文件夹传输
import os
import socket
import threading
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from queue import Queue

class FolderTransfer:
    def __init__(self):
        self.file_queue = Queue()
        self.total_files = 0
        self.transferred_files = 0
        self.total_size = 0
        self.transferred_size = 0
        self.active = True

    def scan_folder(self, folder_path, base_path=""):
        for item in os.listdir(folder_path):
            full_path = os.path.join(folder_path, item)
            rel_path = os.path.join(base_path, item)
            
            if os.path.isdir(full_path):
                self.scan_folder(full_path, rel_path)
            else:
                self.file_queue.put((full_path, rel_path))
                self.total_size += os.path.getsize(full_path)
                self.total_files += 1

    def send_folder(self, conn, folder_path, progress_callback):
        self.scan_folder(folder_path)
        conn.send(f"FOLDER|{self.total_files}|{self.total_size}".encode())
        ack = conn.recv(3)
        
        while not self.file_queue.empty() and self.active:
            full_path, rel_path = self.file_queue.get()
            file_size = os.path.getsize(full_path)
            
            conn.send(f"{rel_path}|{file_size}".encode())
            ack = conn.recv(3)
            
            with open(full_path, 'rb') as f:
                while self.active:
                    data = f.read(4096)
                    if not data: break
                    conn.send(data)
                    self.transferred_size += len(data)
                    progress_callback(self.transferred_size, self.total_size)
            
            self.transferred_files += 1
            progress_callback(self.transferred_files, self.total_files, is_file=True)

    def receive_folder(self, conn, save_path, progress_callback):
        header = conn.recv(1024).decode()
        if header.startswith("FOLDER"):
            _, total_files, total_size = header.split("|")
            self.total_files = int(total_files)
            self.total_size = int(total_size)
            conn.send(b'ACK')
        
        for _ in range(self.total_files):
            if not self.active: break
            
            file_info = conn.recv(1024).decode()
            rel_path, file_size = file_info.split('|')
            file_size = int(file_size)
            
            conn.send(b'ACK')
            full_path = os.path.join(save_path, rel_path)
            os.makedirs(os.path.dirname(full_path), exist_ok=True)
            
            received = 0
            with open(full_path, 'wb') as f:
                while received < file_size and self.active:
                    data = conn.recv(min(4096, file_size - received))
                    if not data: break
                    f.write(data)
                    received += len(data)
                    self.transferred_size += len(data)
                    progress_callback(self.transferred_size, self.total_size)
            
            self.transferred_files += 1
            progress_callback(self.transferred_files, self.total_files, is_file=True)

class FileTransfer:
    def __init__(self):
        self.active = False
        self.current_size = 0
        self.total_size = 0
        self.folder_transfer = FolderTransfer()

    def send_file(self, file_path, port, progress_callback):
        self.active = True
        self.total_size = os.path.getsize(file_path)
        
        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server.bind(('0.0.0.0', port))
        server.listen(1)
        
        conn, addr = server.accept()
        conn.send(f"FILE|{os.path.basename(file_path)}|{self.total_size}".encode())
        ack = conn.recv(3)
        
        self.current_size = 0
        with open(file_path, 'rb') as f:
            while self.active:
                data = f.read(4096)
                if not data: break
                conn.send(data)
                self.current_size += len(data)
                progress_callback(self.current_size, self.total_size)
        
        conn.close()
        server.close()

    def receive_file(self, ip, port, save_path, progress_callback):
        self.active = True
        conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        conn.connect((ip, port))
        
        header = conn.recv(1024).decode()
        if header.startswith("FILE"):
            _, filename, total_size = header.split("|")
            self.total_size = int(total_size)
            save_path = os.path.join(save_path, filename)
            conn.send(b'ACK')
        
        self.current_size = 0
        with open(save_path, 'wb') as f:
            while self.active:
                data = conn.recv(4096)
                if not data: break
                f.write(data)
                self.current_size += len(data)
                progress_callback(self.current_size, self.total_size)
        
        conn.close()

    def stop(self):
        self.active = False
        self.folder_transfer.active = False

class TransferGUI:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("内网极速传文件 v2.0")
        self.transfer = FileTransfer()
        self.setup_ui()
        self.root.protocol("WM_DELETE_WINDOW", self.on_close)
        
    def setup_ui(self):
        self.tabs = ttk.Notebook(self.root)
        self.send_tab = ttk.Frame(self.tabs)
        self.recv_tab = ttk.Frame(self.tabs)
        self.tabs.add(self.send_tab, text="发送文件")
        self.tabs.add(self.recv_tab, text="接收文件")
        self.tabs.pack(expand=1, fill="both")

        # 发送界面
        tk.Label(self.send_tab, text="选择路径:").grid(row=0, column=0, padx=5, pady=5)
        self.send_path = tk.Entry(self.send_tab, width=40)
        self.send_path.grid(row=0, column=1, padx=5, pady=5)
        tk.Button(self.send_tab, text="浏览文件", command=lambda: self.browse_file()).grid(row=0, column=2, padx=5, pady=5)
        tk.Button(self.send_tab, text="浏览文件夹", command=lambda: self.browse_folder()).grid(row=0, column=3, padx=5, pady=5)
        
        tk.Label(self.send_tab, text="本机IP:").grid(row=1, column=0, padx=5, pady=5)
        self.local_ip = tk.Label(self.send_tab, text=self.get_local_ip())
        self.local_ip.grid(row=1, column=1, sticky="w", padx=5, pady=5)
        
        tk.Label(self.send_tab, text="端口:").grid(row=2, column=0, padx=5, pady=5)
        self.port_entry = tk.Entry(self.send_tab, width=10)
        self.port_entry.insert(0, "5000")
        self.port_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5)
        
        self.send_btn = tk.Button(self.send_tab, text="开始发送", command=self.start_send)
        self.send_btn.grid(row=3, column=1, pady=10)
        
        self.send_progress = ttk.Progressbar(self.send_tab, orient="horizontal", length=300, mode="determinate")
        self.send_progress.grid(row=4, column=0, columnspan=4, pady=10)
        self.file_progress = ttk.Progressbar(self.send_tab, orient="horizontal", length=300, mode="determinate")
        self.file_progress.grid(row=5, column=0, columnspan=4, pady=5)
        self.status_label = tk.Label(self.send_tab, text="准备就绪")
        self.status_label.grid(row=6, column=0, columnspan=4)

        # 接收界面
        tk.Label(self.recv_tab, text="发送方IP:").grid(row=0, column=0, padx=5, pady=5)
        self.recv_ip = tk.Entry(self.recv_tab, width=20)
        self.recv_ip.grid(row=0, column=1, sticky="w", padx=5, pady=5)
        
        tk.Label(self.recv_tab, text="保存路径:").grid(row=1, column=0, padx=5, pady=5)
        self.save_path = tk.Entry(self.recv_tab, width=40)
        self.save_path.grid(row=1, column=1, padx=5, pady=5)
        tk.Button(self.recv_tab, text="浏览", command=self.browse_save).grid(row=1, column=2, padx=5, pady=5)
        
        tk.Label(self.recv_tab, text="端口:").grid(row=2, column=0, padx=5, pady=5)
        self.recv_port = tk.Entry(self.recv_tab, width=10)
        self.recv_port.insert(0, "5000")
        self.recv_port.grid(row=2, column=1, sticky="w", padx=5, pady=5)
        
        self.recv_btn = tk.Button(self.recv_tab, text="开始接收", command=self.start_recv)
        self.recv_btn.grid(row=3, column=1, pady=10)
        
        self.recv_progress = ttk.Progressbar(self.recv_tab, orient="horizontal", length=300, mode="determinate")
        self.recv_progress.grid(row=4, column=0, columnspan=3, pady=10)
        self.recv_file_progress = ttk.Progressbar(self.recv_tab, orient="horizontal", length=300, mode="determinate")
        self.recv_file_progress.grid(row=5, column=0, columnspan=3, pady=5)
        self.recv_status = tk.Label(self.recv_tab, text="等待连接...")
        self.recv_status.grid(row=6, column=0, columnspan=3)

    def browse_file(self):
        path = filedialog.askopenfilename()
        if path:
            self.send_path.delete(0, tk.END)
            self.send_path.insert(0, path)
            self.send_progress["maximum"] = os.path.getsize(path)

    def browse_folder(self):
        path = filedialog.askdirectory()
        if path:
            self.send_path.delete(0, tk.END)
            self.send_path.insert(0, path)
            total_size = sum(os.path.getsize(os.path.join(dirpath, filename)) 
                            for dirpath, _, filenames in os.walk(path) 
                            for filename in filenames)
            self.send_progress["maximum"] = total_size

    def browse_save(self):
        path = filedialog.asksaveasfilename()
        if not path: return
        self.save_path.delete(0, tk.END)
        self.save_path.insert(0, path)

    def start_send(self):
        path = self.send_path.get()
        if not path:
            messagebox.showerror("错误", "请选择文件或文件夹!")
            return
            
        port = int(self.port_entry.get())
        self.send_btn.config(state="disabled")
        
        if os.path.isdir(path):
            threading.Thread(target=self.transfer.folder_transfer.send_folder,
                           args=(self.create_server(port), path, self.update_send_progress),
                           daemon=True).start()
            self.status_label.config(text="等待接收方连接文件夹...")
        else:
            threading.Thread(target=self.transfer.send_file,
                           args=(path, port, self.update_send_progress),
                           daemon=True).start()
            self.status_label.config(text="等待接收方连接文件...")

    def start_recv(self):
        if not self.save_path.get():
            messagebox.showerror("错误", "请选择保存路径!")
            return
            
        ip = self.recv_ip.get()
        port = int(self.recv_port.get())
        self.recv_btn.config(state="disabled")
        
        threading.Thread(target=self.handle_receive,
                        args=(ip, port, self.save_path.get()),
                        daemon=True).start()
        self.recv_status.config(text="正在连接发送方...")

    def handle_receive(self, ip, port, save_path):
        conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        conn.connect((ip, port))
        
        header = conn.recv(1024).decode()
        conn.send(b'ACK')
        
        if header.startswith("FOLDER"):
            if not os.path.exists(save_path):
                os.makedirs(save_path)
            self.transfer.folder_transfer.receive_folder(conn, save_path, self.update_recv_progress)
        else:
            self.transfer.receive_file(ip, port, save_path, self.update_recv_progress)
        
        conn.close()

    def create_server(self, port):
        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server.bind(('0.0.0.0', port))
        server.listen(1)
        conn, _ = server.accept()
        return conn

    def update_send_progress(self, current, total, is_file=False):
        if is_file:
            self.file_progress["value"] = current
            self.status_label.config(text=f"文件进度: {current}/{total}")
        else:
            percent = (current/total)*100 if total > 0 else 0
            self.send_progress["value"] = current
            self.status_label.config(text=f"总进度: {current/1024/1024:.1f}MB/{total/1024/1024:.1f}MB ({percent:.1f}%)")

    def update_recv_progress(self, current, total, is_file=False):
        if is_file:
            self.recv_file_progress["maximum"] = total
            self.recv_file_progress["value"] = current
            self.recv_status.config(text=f"文件进度: {current}/{total}")
        else:
            self.recv_progress["maximum"] = total
            self.recv_progress["value"] = current
            percent = (current/total)*100 if total > 0 else 0
            self.recv_status.config(text=f"总进度: {current/1024/1024:.1f}MB/{total/1024/1024:.1f}MB ({percent:.1f}%)")

    def get_local_ip(self):
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            ip = s.getsockname()[0]
            s.close()
            return ip
        except:
            return "127.0.0.1"

    def on_close(self):
        self.transfer.stop()
        self.root.destroy()

    def run(self):
        self.root.mainloop()

if __name__ == "__main__":
    app = TransferGUI()
    app.run()
